import torch
from PreRec.BPRMF.Trainer import Trainer
from PreRec.BPRMF.Model import MF
from PreRec.BPRMF.FileLoader import FileLoader
import time
from tqdm import tqdm
from typing import List, Dict

from utils.log_output import simple_text_log


class Recommender:
    def __init__(self, hyper_params):
        self.hyper_params = hyper_params
        self.file_loader = FileLoader(hyper_params)
        self.mf_model = MF(hyper_params)

        self.mf_model = self.mf_model.to(hyper_params['device'])
        trainer = Trainer(hyper_params)
        trainer.mf_model = self.mf_model
        trainer.fit(save_model=True)

        self.item_embed = self.mf_model.item_embedding.weight
        self.user_embed = self.mf_model.user_embedding.weight

    def recommend(self, rec_count=500) -> Dict[int, List[int]]:
        print('BPRMF: Start BPRMF recommend...')
        simple_text_log('train', 'Start BPR-MF recommendation...')
        item_cnt = self.hyper_params['item_cnt']
        user_cnt = self.hyper_params['user_cnt']

        self.mf_model.eval()
        self.coarse_result = torch.empty(
            [user_cnt, rec_count], dtype=torch.long)

        batch_size = self.hyper_params['batch_size']
        batch_cnt = user_cnt // batch_size

        with torch.no_grad():
            for i in tqdm(range(batch_cnt)):
                left = i * batch_size
                right = left + batch_size
                if i == batch_cnt - 1:
                    right = user_cnt

                cur_user = self.user_embed[left:right].unsqueeze(1).expand(
                    [-1, item_cnt, -1])
                mf_result = torch.sum(cur_user*self.item_embed, dim=-1)
                self.coarse_result[left:right] = torch.argsort(
                    mf_result, dim=-1, descending=True)[:, :rec_count]

        result: Dict[int, List[int]] = {}
        for user in range(user_cnt):
            result[user] = self.coarse_result[user].tolist()

        return result


if __name__ == '__main__':
    hyper_params = {
        'dataset_path': '../datasets/lib.txt',
        'device': 'cpu',
        'batch_size': 128,
        'epochs': 2,
        'embed_dim': 32,
        'neg_sample_cnt': 500,
        'epsilon': 0.1
    }

    recommender = Recommender(hyper_params)

    start_time = time.time()
    result = recommender.recommend()
    print(f'BPRMF Time: {time.time() - start_time}')

    print(result)
